from bean.product import Product
from bean.equipment import Equipment
from bean.route import Route


class Data:
    def __init__(self, product_path="data/产品信息.csv",
                 equip_path="data/设备信息.csv",
                 route_path="data/工艺路线.csv"):
        self.products = []
        self.equips = []
        self.routes = []
        with open(product_path, 'r', encoding='utf-8') as f:
            for line in f.readlines()[1:]:
                self.products.append(Product(*line.strip().split(',')))
        self.products.sort(key=lambda x: (x.route_id, x.section_num))

        with open(equip_path, 'r', encoding='gbk') as f:
            for line in f.readlines()[1:]:
                self.equips.append(Equipment(*line.strip().split(',')))

        route_d = {}
        with open(route_path, 'r', encoding='gbk') as f:
            for line in f.readlines()[1:]:
                line = line.strip().split(',')
                if line[0] not in route_d:
                    route_d[line[0]] = []
                route_d[line[0]].append(line[1:])

        for k, v in route_d.items():
            self.routes.append(
                Route(
                    k,
                    [x[0] for x in v],
                    [x[1] for x in v],
                    [x[2] for x in v],
                    [x[3].strip().split('-')[0] for x in v],
                    [x[3].strip().split('-')[-1] for x in v],
                    [x[4] for x in v],
                    [x[5] for x in v],
                )
            )

    def show_data(self):
        print("product info: ")
        for p in self.products:
            p.output()
        print("equip info: ")
        for equip in self.equips:
            equip.output()
        print("route info: ")
        for route in self.routes:
            route.output()

    def write_ans(self, product_dispatchs, ans_path="../ans/ans.csv"):
        with open(ans_path, "w", encoding="utf-8") as f:
            f.write("product_id,route_No,equ_name,start,duration,end\n")
            for dispatch in product_dispatchs:
                for row in dispatch.write_ans():
                    f.write(row + "\n")


if __name__ == '__main__':
    data = Data()
    data.show_data()
